import pandas as pd
import os
import numpy as np
from sklearn.model_selection import train_test_split


class AdultDataset(object):

    def __init__(self, data_path="data/adult"):
        cols = [
            "age",
            "workclass",
            "fnlwgt",
            "education",
            "education_num",
            "marital_status",
            "occupation",
            "relationship",
            "race",
            "sex",
            "capital_gain",
            "capital_loss",
            "hours_per-week",
            "native_country",
        ]

        categorical_cols = [
            "workclass",
            "education",
            "marital_status",
            "occupation",
            "relationship",
            "race",
            "sex",
            "native_country",
        ]

        train_df = pd.read_csv(os.path.join(data_path, "adult.data"), skiprows=1, names=cols + ['target'], na_values="?")
        test_df = pd.read_csv(os.path.join(data_path, "adult.test"), skiprows=1, names=cols + ['target'], na_values="?")

        all_df = pd.concat([train_df, test_df], axis=0)
        all_df = pd.get_dummies(all_df, columns=categorical_cols + ['target'])
        train_df = all_df.iloc[:train_df.shape[0], :]
        test_df = all_df.iloc[train_df.shape[0]:, :]
        self.train_X = train_df.iloc[:, :-4].values
        self.train_y = train_df['target_ <=50K'].values
        self.test_X = test_df.iloc[:, :-4].values
        self.test_y = test_df['target_ <=50K.'].values

class CaliforniaHousingDataset(object):

    def __init__(self, data_path='data/CaliforniaHousing/', test_size=0.2, random_state=None):
        self.df = pd.read_csv(os.path.join(data_path, 'cal_housing.data'), header=-1,
            names=['f1','f2','f3','f4','f5','f6','f7','f8','target'])

        self.X = self.df.filter(axis=1, regex="f.*").values

        self.y = self.df['target'].values

        self.y = np.log(self.y)

        self.train_X, self.test_X, self.train_y, self.test_y = \
            train_test_split(self.X, self.y, test_size=0.2, random_state=None)